#!/bin/bash

### Functions ###


############################### DISK FUNCTIONS START ##########################
## Start fs functions

#
# Desc: fjerner alt av partisjoner på usb pennen
function deleteallpart() {
# Funksjonen ser ut til å være ustabil ennå...
    echo "[+] Nulling out all filesystems on $USB_DEVICE"
    dd count=1 bs=1024 if=/dev/zero of=$USB_DEVICE || failure "Could not delete contents.."
}

#
# Desc: lager et dos partisjonssystem
function make_dos_partition() {
    deleteallpart
    # Lager 2x partisjoner:
    # Part01, 800Mb size, FAT16 filesystem
    # Part02, Rest size, Linux Ext2
    sfdisk -uM $USB_DEVICE <<EOF
,800,6,*
,,83
EOF
}

#
# Desc: lager et gpt partisjonssystem
function make_mac_partition() {
    deleteallpart
    # Lager 2x partisjoner:
    # Part01, 800Mb size, FAT16 filesystem
    # Part02, Rest size, Linux Ext2
    parted -s $USB_DEVICE mklabel gpt
    parted -s $USB_DEVICE mklabel primary 1 840
    parted -s $USB_DEVICE mkpartfs primary ext2 841 1002
}

#
# Desc: Lager en 800mb hfsplus partisjon og en ext2 til resterende
function format_mac_dev() {
    echo "[+] Formatting the usb device $USB_DEVICE"
    mkfs.hfsplus -v digeksUSB $USB_PART1 || failure "Could not format partition $USB_PART1"
    mkfs.ext2 -b 4096 -L casper-rw $USB_PART2 || failure "Could not format partititon $USB_PART2"
    parted -s $USB_DEVICE set 1 boot on
    gptsync $USB_DEVICE || failure "Failed to sync partition table on $USB_DEVICE"
}

#
# Desc: Lager en 800mb fat16 partisjon og en ext2 til resterende
function format_dos_dev() {
    echo "[+] Formatting the usb device $USB_DEVICE"
    mkfs.vfat -F 16 -n digeksUSB $USB_PART1 || failure "Could not format partition $USB_PART1"
    mkfs.ext2 -b 4096 -L casper-rw $USB_PART2 || failure "Could not format partition $USB_PART2"
}


#
# Desc: installerer mbr softwaren på pc'en om den ikke har det..
function check_mbr_software() {
    if [ -x /sbin/install-mbr ] ; then
        echo "[+] install-mbr found, installing mbr to usb..."
    else
        echo "[-] mbr not installed, trying to install it.."
        # Installing mbr software
        apt-get -y --force-yes install mbr
        if [ -x /sbin/install-mbr ] ; then
            echo "[+] Install complete.."
        else
            echo "[-] install-mbr still not found, exiting.."
            exit 1
        fi
    fi
}


#
# Desc: installerer syslinux bootloader
function bootable_dos() {
    echo "[+] Unmounting and making the usb partition bootable"
    # MS-DOS
    cd $MNT_FOLDER
#    umount $MNT_FOLDER || failure "Could not umount $USB_PART1"
    sync
    syslinux -d syslinux -f $USB_PART1 || failure "Could not make it bootable"
    cd $DIGEKS_SOURCE
    check_mbr_software
    echo "[+] Installing MBR to $USB_DEVICE"
    install-mbr $USB_DEVICE || failure "Could not set mbr"
    sync
}

#
# Desc: installerer grub2 med efi support
# UNSUPPORTED FUNCTION!!!
function bootable_mac() {
    echo "[+] Making the liveusb bootable on mac"
    GRUBDIR=$MNT_FOLDER/efi/grub
    GRUBDIR2=$MNT_FOLDER/Boot/grub
    LIBDIR=/usr/lib/grub/i386-efi
    mkdir -p $GRUBDIR
    mkdir -p $GRUBDIR2
    grub-mkimage -o $GRUBDIR/grub.efi gpt hfsplus fat
    cp $LIBDIR/*.mod $LIBDIR/fs.lst $LIBDIR/command.lst $GRUBDIR
    cp $CUSTOMIZE_DIR/menu.lst $GRUBDIR
    cp $GRUBDIR/* $GRUBDIR2
    #echo "Grub-install ..."
    #grub-install --grub-mkimage=$GRUBDIR/grub.efi --root-directory=$USB_MOUNT --recheck
    sync
}

#
# Desc: Gjør usb devicen bootbar, alias funksjon..
function make_bootable() {
#    if [ "$FS_TYPE" == 1 ] ; then
#        bootable_dos
#    else
#        # Mac / GPT
#        bootable_mac
#    fi
    bootable_dos
}
## End fs functions

#
# Desc: Avmonterer alt som har med usb device'n å gjør..
function umountdev() {
    ismounted=$(mount |grep $USB_DEVICE)
    if [[ ${#ismounted} > 0 ]]; then
        echo "[+] Unmounting the device $USB_DEVICE"
        umount $USB_PART1
        umount $USB_PART2
    fi
}

#
# Desc: Reset på mountpoint
function remkmp() { # Remove temp mount point
    rm -fr $MNT_FOLDER
    mkdir -p $MNT_FOLDER
}

#
# Desc: mount fat16 elr hfsplus partisjonen..
function mountdev() {
    echo "[+] Mounting usb partition $USB_PART1 at $MNT_FOLDER"
    remkmp
    mount $USB_PART1 $MNT_FOLDER
}

############################### DISK FUNCTIONS STOP ##########################


#
# Desc: forklarer hvordan du kan bruke scriptet
function usage() {
    echo ""
    echo ""
    echo "[+] WRiteUSBIMaGe Usage menu: "
    echo "  -d [device-name] (optimal if digeksconfig is present with usbdevice & path prefix) "
    echo "  -i [image-folder] (optimal if digeksconfig is present with usbdevice & path prefix) "
    echo "  -c [config-file] (optimal) "
    echo "  -f   Force, ignore digeksconfig "
    echo "  -m   MSDOS Partition Table DEFAULT=YES "
    echo "  -x   GPT Partition Table. UNSUPPORTED ATM! "
    echo ""
    exit 1
}

#
# Desc: feilmelding..
function failure() {
    echo "[-] $@"
    exit 1
}


### Check if it's missing parameters or you're not root... ###
if [ -z $1 ]; then
    usage
fi

if [ ! `id -un` = "root" ] ; then
    echo "[-] $0 should be run as root! Use at least sudo...."
    exit 1
fi


### Checking parameters ###
while getopts "d:i:m:x:h:f:c:" options; do
    case $options in
        m ) export FS_TYPE=1;;
        x ) export FS_TYPE=2;;
        f ) export ENABLE_FORCE=1;;
        c ) export CONFIG_FILE=$OPTARG;;
        d ) export OPT_USB_DEVICE=$OPTARG;;
        i ) export OPT_IMG_FOLDER=$OPTARG;;
        h ) usage;;
    esac
done


# Funka ikke...
echo "[+] Reading config"
if [ -z "$CONFIG_FILE" ] ; then
    export CONFIG_FILE=$(pwd)"/digeksconfig"
    echo "[+] Config: digeksconfig FOUND!"
    if [ -d $MNT_FOLDER ] ; then
        rm -fr $MNT_FOLDER
        mkdir -p $MNT_FOLDER
    else
        mkdir -p $MNT_FOLDER
    fi
fi

if [ -n $ENABLE_FORCE ] ; then
    echo "[+] Force enabled... Ignoring the config file.."
    export USB_DEVICE=$OPT_USB_DEVICE
    export IMG_FOLDER=$OPT_IMG_FOLDER
else
    ### Checking variables (and config file) ###
    echo "[-] Missing parameters, checking if digeksconfig has the information needed."
    source $CONFIG_FILE
    if [ -o "$PREDEF_USB_DEVICE" ] ; then
        export USB_DEVICE=$PREDEF_USB_DEVICE
    fi
    if [ -o "$PREDEF_MNT_FOLDER" ] ; then
        export MNT_FOLDER=$PREDEF_MNT_FOLDER
    fi
    if [ -z $USB_DEVICE ] ; then
        echo "[-] Missing USB_DEVICE, please add it to digeksconfig or use $0 -d [device-name](e.g. /dev/sdb)"
        usage
    fi
    if [ -z $IMG_FOLDER ] ; then
        echo "[-] Missing IMG_FOLDER/RESULT_FOLDER, please add it to digeksconfig or use $0 -i [path-name]"
        usage
    fi
fi

# End checking, finaly work...

##############


#
# Desc: denne funksjonen flytter kjernen
#function move_kernel() {
#    echo "[+] Moving kernel and initrd to root of usb"
#    cd $IMG_FOLDER
#    cp -f $PRE_FOLDER/casper/vmlinuz $MNT_FOLDER
#    cp -f $PRE_FOLDER/casper/initrd.gz $MNT_FOLDER
# Det er noen filer i isolinux som gjør at boot henger med "kernel panic" (muligens isolinux.bin)
#    cd $PRE_FOLDER
#    sync
#}

#
# Desc: Kopierer selve operativ systemet over til usb disken
function copy_files() {
    echo "[+] Copy files from $IMG_FOLDER to $MNT_FOLDER. Takes a few minutes."
    #(cd $IMG_FOLDER ; tar -cf - *) | (cd $MNT_FOLDER ; tar -xvf - )
    # copy command produces some error messages about
    # "cannot create symbolic link..." and "failed to preserve ownership..."
    # Ignore these

    cp -rv $IMG_FOLDER/* $MNT_FOLDER
    sync
}

#
# Desc: alt i en funksjon, tar seg av prosessen
function doitall() {
    mountdev
    cd $DIGEKS_SOURCE
    echo "[+] Copying files to usb disk.."
    copy_files
    #echo "[+] Copying kernel to usb disk..."
   # move_kernel
    echo "DEBUG: make_bootable"
    make_bootable
    echo "DEBUG: umountdev"
    umountdev
}




# Setting last variables
export USB_PART1=${USB_DEVICE}1
export USB_PART2=${USB_DEVICE}2

if [ -z $MNT_FOLDER ] ; then
    export MNT_FOLDER=`pwd`/result/mount/
fi
if [ -z $PRE_FOLDER ] ; then
    export PRE_FOLDER=`pwd`/diskimage/
fi

cd $DIGEKS_SOURCE

## Starter med filsystem/partisjon greia
umountdev
remkmp
deleteallpart

#if [ "$FS_TYPE" == 1 ] ; then
#    # MS-DOS
#    make_dos_partition
#    format_dos_dev
#    doitall
#else # OSX / GPT
#    make_mac_partition
#    format_mac_dev
#    doitall
#fi

echo "DEBUG: make_dos_partition"
make_dos_partition
echo "DEBUG: format_dos_dev"
format_dos_dev
doitall


echo "[+] Process complete... USB Disk should be finished now..."



# EOF
